"""Example Algorithm of AutoEncoder."""

import torch
from torch import nn
from torchvision import datasets, transforms
import matplotlib.pyplot as plt


class AutoEncoder:
    def __init__(self):
        self.encoder = nn.Sequential()
        self.encoder.append(nn.Linear(input_size, hidden_size))
        self.encoder.append(nn.ReLU())
        self.encoder.append(nn.Linear(hidden_size, latent_size))

        self.decoder = nn.Sequential()
        self.decoder.append(nn.Linear(latent_size, hidden_size))
        self.decoder.append(nn.ReLU())
        self.decoder.append(nn.Linear(hidden_size, output_size))
        self.decoder.append(nn.Sigmoid())

        self.params = [p for p in self.encoder.parameters()] + [p for p in self.decoder.parameters()]
        self.optimizer = torch.optim.Adam(self.params, lr=learning_rate)

    def update(self, state):
        feature = self.encoder(state)
        rebuild = self.decoder(feature)
        loss = torch.nn.BCELoss(reduction='sum')(rebuild, state)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        return loss

    def save_model(self, filename="autoencoder.pt"):
        state_dict = {"encoder": self.encoder.state_dict(), "decoder": self.decoder.state_dict()}
        torch.save(state_dict, filename)

    def load_model(self, filename="autoencoder.pt"):
        state_dict = torch.load(filename)
        self.encoder.load_state_dict(state_dict["encoder"])
        self.decoder.load_state_dict(state_dict["decoder"])

    def activate(self, state):
        feature = self.encoder(state)
        # feature = torch.tensor([0 for _ in range(16)], dtype=torch.float32)
        img = self.decoder(feature)
        plt.matshow(img.view(28, 28).cpu().detach().numpy())
        plt.show()


def train():
    train_data = datasets.MNIST('/data', train=True, transform=transforms.ToTensor())
    train_loader = torch.utils.data.DataLoader(train_data, batch_size=32, shuffle=True)
    for epoch in range(epochs):
        for imgs, _ in train_loader:
            imgs = imgs.view(imgs.shape[0], input_size)  # imgs:(bs,28*28)
            agent.update(imgs)
    agent.save_model()


def test():
    agent.load_model()
    test_data = datasets.MNIST('/data', train=False, transform=transforms.ToTensor())
    test_loader = torch.utils.data.DataLoader(test_data, batch_size=32, shuffle=False)
    for imgs, _ in test_loader:
        imgs = imgs.view(imgs.shape[0], input_size)  # imgs:(bs,28*28)
        agent.activate(imgs[0])
        break


latent_size = 16
hidden_size = 128
input_size = output_size = 28 * 28

epochs = 20
batch_size = 32
learning_rate = 1e-3

agent = AutoEncoder()
